package io.wen.bd.s6m3.spark

import java.util.concurrent.Future

import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord, RecordMetadata}

class SparkKafkaProducer[K, V](createProducer: () => KafkaProducer[K, V]) extends Serializable {
  // 关键所在，使用lazy变量，将producer实例化延迟到excutors的第一次使用，避免了未序列化异常
  lazy val producer = createProducer()

  def send(topic: String, key: K, value: V): Future[RecordMetadata] =
    producer.send(new ProducerRecord[K, V](topic, key, value))

  def send(topic: String, value: V): Future[RecordMetadata] =
    producer.send(new ProducerRecord[K, V](topic, value))
}

object SparkKafkaProducer {
  import scala.collection.JavaConversions._

  def apply[K, V](config: Map[String, Object]): SparkKafkaProducer[K, V] = {
    val createTopicFunc = () => {
      val producer = new KafkaProducer[K, V](config)
      sys.addShutdownHook(producer.close())
      producer
    }
    new SparkKafkaProducer(createTopicFunc)
  }

  def apply[K, V](config: java.util.Properties): SparkKafkaProducer[K, V] = apply(config.toMap)
}
